import numpy as np
from skfuzzy import control as ctrl
from skfuzzy import membership
import matplotlib.pyplot as plt

# 定义ANFIS结构
class ANFIS:
    def __init__(self):
        # 定义输入和输出变量
        self.input1 = ctrl.Antecedent(np.arange(0, 10, 1), 'input1')
        self.input2 = ctrl.Antecedent(np.arange(0, 10, 1), 'input2')
        self.output = ctrl.Consequent(np.arange(0, 100, 1), 'output')

        # 定义模糊集合
        self.input1.automf(3)
        self.input2.automf(3)
        self.output['low'] = membership.trimf(self.output.universe, [0, 0, 50])
        self.output['medium'] = membership.trimf(self.output.universe, [0, 50, 100])
        self.output['high'] = membership.trimf(self.output.universe, [50, 100, 100])

        # 定义规则
        rule1 = ctrl.Rule(self.input1['poor'] & self.input2['poor'], self.output['low'])
        rule2 = ctrl.Rule(self.input1['average'] & self.input2['average'], self.output['medium'])
        rule3 = ctrl.Rule(self.input1['good'] & self.input2['good'], self.output['high'])

        self.control_system = ctrl.ControlSystem([rule1, rule2, rule3])
        self.simulator = ctrl.ControlSystemSimulation(self.control_system)

    def evaluate(self, x1, x2):
        self.simulator.input['input1'] = x1
        self.simulator.input['input2'] = x2
        self.simulator.compute()
        return self.simulator.output['output']

# 定义卷积操作
def convolve_with_anfis(image, anfis_model, stride=1, padding=0):
    # 图像填充
    padded_image = np.pad(image, padding, mode='constant')
    output_shape = ((padded_image.shape[0] - 2) // stride + 1,
                    (padded_image.shape[1] - 2) // stride + 1)
    output_image = np.zeros(output_shape)

    for i in range(0, padded_image.shape[0] - 2 + 1, stride):
        for j in range(0, padded_image.shape[1] - 2 + 1, stride):
            region = padded_image[i:i + 3, j:j + 3]
            # 使用ANFIS处理3x3区域
            output_image[i // stride, j // stride] = anfis_model.evaluate(np.mean(region), np.std(region))

    return output_image

# 测试代码
if __name__ == '__main__':
    # 创建ANFIS模型
    anfis_model = ANFIS()

    # 测试图像
    test_image = np.random.rand(5, 5) * 10  # 5x5的随机图像

    # 卷积
    result_image = convolve_with_anfis(test_image, anfis_model, stride=1, padding=1)

    # 可视化结果
    plt.figure(figsize=(8, 4))
    plt.subplot(1, 2, 1)
    plt.title('原始图像')
    plt.imshow(test_image, cmap='gray', vmin=0, vmax=10)
    plt.colorbar()

    plt.subplot(1, 2, 2)
    plt.title('卷积后的图像')
    plt.imshow(result_image, cmap='gray', vmin=0, vmax=100)
    plt.colorbar()

    plt.show()
